package tt.langchain4j.store;

import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.*;
import dev.langchain4j.store.embedding.filter.Filter;
import org.redisson.api.RMap;
import org.redisson.api.RSet;
import org.redisson.api.RedissonClient;
import org.redisson.client.codec.StringCodec;
import org.redisson.codec.JsonJacksonCodec;

import java.util.*;

import static dev.langchain4j.internal.Utils.randomUUID;
import static dev.langchain4j.internal.ValidationUtils.ensureNotNull;
import static java.util.Comparator.comparingDouble;

public class RedissonEmbeddingStore implements EmbeddingStore<TextSegment> {

    private final RedissonClient redissonClient;

    private final Map<String,Entry> memoryEntry = new HashMap<>();

    private final RMap<String, CopyEmbedding> embedding;
    private final RMap<String, CopyTextSegment> embedded;
    private final RSet<String> ids;



    public static class Entry {

        private String id;

        private Embedding embedding;

        private TextSegment textSegment;

        public static Entry from( String id , Embedding embedding , TextSegment textSegment ) {
            final Entry entry = new Entry();
            entry.id = id;
            entry.embedding = embedding;
            entry.textSegment = textSegment;
            return entry;
        }
    }




    public RedissonEmbeddingStore(RedissonClient redissonClient ) {
        this.redissonClient = redissonClient;
        //初始化
        this.embedding = redissonClient.getMap("Embedding" , new JsonJacksonCodec());
        this.embedded = redissonClient.getMap("Embedded",new JsonJacksonCodec());
        this.ids = redissonClient.getSet("Ids",new StringCodec());
    }




    @Override
    public String add(Embedding embedding) {
        String id = randomUUID();
        add(id, embedding);
        return id;
    }

    @Override
    public void add(String id, Embedding embedding) {
        add(id, embedding, null);
    }

    @Override
    public String add(Embedding embedding, TextSegment embedded) {
        String id = randomUUID();
        add(id, embedding, embedded);
        return id;
    }

    public void add(String id, Embedding embedding, TextSegment embedded) {
        //写入到内存
        memoryEntry.put( id , Entry.from(id,embedding,embedded) );
        //写入到redis中
        this.ids.add(id);
        this.embedding.put(id, CopyEmbedding.from(embedding));
        this.embedded.put(id,CopyTextSegment.from(embedded));
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings) {
        List<String> rs = new ArrayList<>(embeddings.size());
        for (Embedding _embedding : embeddings) {
            rs.add(add(_embedding));
        }
        return rs;
    }

    @Override
    public void addAll(List<String> ids, List<Embedding> embeddings, List<TextSegment> embedded) {
        if (ids.size()!= embeddings.size() || embeddings.size() != embedded.size()) {
            throw new IllegalArgumentException("The list of ids and embeddings and embedded must have the same size");
        }
        for (int i = 0; i < ids.size(); i++) {
            add(ids.get(i),embeddings.get(i),embedded.get(i));
        }
    }

    @Override
    public void removeAll(Collection<String> ids) {
        this.ids.removeAll(ids);
        for (String id : ids) {
            this.memoryEntry.remove(id);
            this.embedding.remove(id);
            this.embedded.remove(id);
        }
    }

    @Override
    public void removeAll(Filter filter) {
        ensureNotNull(filter, "filter");
        //TODO redis 清除


    }

    @Override
    public void removeAll() {
        ids.clear();
        embedding.clear();
        embedded.clear();
    }

    @Override
    public EmbeddingSearchResult<TextSegment> search(EmbeddingSearchRequest embeddingSearchRequest) {
        Comparator<EmbeddingMatch<TextSegment>> comparator = comparingDouble(EmbeddingMatch::score);
        PriorityQueue<EmbeddingMatch<TextSegment>> matches = new PriorityQueue<>(comparator);
        Filter filter = embeddingSearchRequest.filter();
        for (String _id : ids) {
            Entry entry = memoryEntry.compute(_id, (k, v) -> {
                if ( null == v ) {
                    //初始化
                    v = new Entry();
                    v.id = k;
                    v.embedding = CopyEmbedding.to(RedissonEmbeddingStore.this.embedding.get(k));
                    v.textSegment = CopyTextSegment.to( RedissonEmbeddingStore.this.embedded.get(k) );
                }
                return v;
            });
            if (filter != null &&  entry.embedding != null  ) {
                Metadata metadata = ((TextSegment) embedded).metadata();
                if (!filter.test(metadata)) {
                    continue;
                }
            }
            if ( null != entry.textSegment && null != entry.embedding ) {
                double cosineSimilarity = CosineSimilarity.between(entry.embedding, embeddingSearchRequest.queryEmbedding());
                double score = RelevanceScore.fromCosineSimilarity(cosineSimilarity);
                if (score >= embeddingSearchRequest.minScore()) {
                    matches.add(new EmbeddingMatch<>(score, _id, entry.embedding, entry.textSegment));
                    if (matches.size() > embeddingSearchRequest.maxResults()) {
                        matches.poll();
                    }
                }
            }
        }

        List<EmbeddingMatch<TextSegment>> result = new ArrayList<>(matches);
        result.sort(comparator);
        Collections.reverse(result);
        return new EmbeddingSearchResult<>(result);
    }


    public static class CopyEmbedding {

        private float[] vector;

        public static CopyEmbedding from(Embedding embedding  ) {
            if ( null == embedding ) {
                return null;
            }
            CopyEmbedding copyEmbedding = new CopyEmbedding();
            copyEmbedding.vector = embedding.vector();
            return copyEmbedding;
        }

        public static Embedding to( CopyEmbedding copyEmbedding ) {
            if ( null == copyEmbedding || null == copyEmbedding.vector ) {
                return null;
            }
            return new Embedding(copyEmbedding.vector);
        }

        public float[] getVector() {
            return vector;
        }

        public void setVector(float[] vector) {
            this.vector = vector;
        }
    }

    public static class CopyTextSegment {

        private String text;

        private Map<String, Object> metadata;

        public static CopyTextSegment from( TextSegment textSegment ) {
            if ( null == textSegment ) {
                return null;
            }
            CopyTextSegment copyTextSegment = new CopyTextSegment();
            copyTextSegment.text = textSegment.text();
            copyTextSegment.metadata = textSegment.metadata().toMap();
            return copyTextSegment;
        }

        public static TextSegment to( CopyTextSegment textSegment ) {
            if (  null == textSegment || null == textSegment.text || textSegment.text.isBlank() ) {
                return null;
            }
            return new TextSegment( textSegment.text , new Metadata(textSegment.metadata) );
        }

        public String getText() {
            return text;
        }

        public void setText(String text) {
            this.text = text;
        }

        public Map<String, Object> getMetadata() {
            return metadata;
        }

        public void setMetadata(Map<String, Object> metadata) {
            this.metadata = metadata;
        }
    }

    public static class CopyEntry {

        private String id;

        private CopyEmbedding embedding;

        private CopyTextSegment embedded;

        public String getId() {
            return id;
        }

        public void setId(String id) {
            this.id = id;
        }

        public CopyEmbedding getEmbedding() {
            return embedding;
        }

        public void setEmbedding(CopyEmbedding embedding) {
            this.embedding = embedding;
        }

        public CopyTextSegment getEmbedded() {
            return embedded;
        }

        public void setEmbedded(CopyTextSegment embedded) {
            this.embedded = embedded;
        }
    }
}
